package org.lzfto.flink.demo.business.flinkTest.source.mqtt;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.apache.flink.api.connector.source.Boundedness;
import org.apache.flink.api.connector.source.Source;
import org.apache.flink.api.connector.source.SourceReader;
import org.apache.flink.api.connector.source.SourceReaderContext;
import org.apache.flink.api.connector.source.SplitEnumerator;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.core.io.SimpleVersionedSerializer;

/**
 * mqtt数据源-读取特定主题的数据，并写消息到特定的websocket客户端
 * </br> 注意，本source不支持并发，但是可以创建多个实例，多个实例要求mqttClientId不同
 * </br> 构建函数所需要参数，参见： {@linkplain org.lzfto.flink.demo.business.flinkTest.source.mqtt.MqttConnectionConfig 配置 }
 */
public class MqttSource implements Source<String, MqttSplit, MqttEnumeratorCheckpoint> {
    private static final long serialVersionUID = 1L;

    private MqttConnectionConfig config;

    public MqttSource(MqttConnectionConfig config) {
        this.config=config;
    }

    @Override
    public Boundedness getBoundedness() {
        return Boundedness.CONTINUOUS_UNBOUNDED;
    }

    @Override
    public SourceReader<String, MqttSplit> createReader(SourceReaderContext readerContext) throws Exception {
    	return new MqttSourceReader(config);
    }
    

    @Override
    public SimpleVersionedSerializer<MqttSplit> getSplitSerializer() {
        return new MqttSplitSerializer();
    }

    

    private static class MqttSplitSerializer implements SimpleVersionedSerializer<MqttSplit> {
        @Override
        public int getVersion() {
            return 1;
        }

        @Override
        public byte[] serialize(MqttSplit split) {
            return (split.splitId() + ";" + split.getTopic() + ";" + split.getQos() + ";" + split.getClientId()).getBytes();
        }

        @Override
        public MqttSplit deserialize(int version, byte[] serialized) {
            String[] parts = new String(serialized).split(";");
            return new MqttSplit(parts[0], parts[1], Integer.parseInt(parts[2]), parts[3]);
        }
    }

	@Override
	public SplitEnumerator<MqttSplit, MqttEnumeratorCheckpoint> createEnumerator(
			SplitEnumeratorContext<MqttSplit> enumContext) throws Exception {
		List<MqttSplit> initialSplits = new ArrayList<>();
		initialSplits.add(new MqttSplit("split-0", config.getTopic(), config.getQos(), config.getMqttClientId()));
		return new MqttEnumerator(enumContext, initialSplits);
	}

	@Override
	public SplitEnumerator<MqttSplit, MqttEnumeratorCheckpoint> restoreEnumerator(
			SplitEnumeratorContext<MqttSplit> enumContext, MqttEnumeratorCheckpoint checkpoint) throws Exception {
		List<MqttSplit> splits = new ArrayList<>();
		for (byte[] serializedSplit : checkpoint.getSerializedAllSplits()) {
			splits.add(getSplitSerializer().deserialize(1, serializedSplit));
		}
		MqttEnumerator enumerator = new MqttEnumerator(enumContext, splits);
		
		// 恢复待分配的分片
		for (Map.Entry<Integer, List<byte[]>> entry : checkpoint.getSerializedPendingSplits().entrySet()) {
			List<MqttSplit> pendingSplits = new ArrayList<>();
			for (byte[] serializedSplit : entry.getValue()) {
				pendingSplits.add(getSplitSerializer().deserialize(1, serializedSplit));
			}
			enumerator.addSplitsBack(pendingSplits, entry.getKey());
		}
		
		return enumerator;
	}

	@Override
	public SimpleVersionedSerializer<MqttEnumeratorCheckpoint> getEnumeratorCheckpointSerializer() {
		return new JsonSimpleVersionedSerializer();
	}
}
	